{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using LlamaIndex-Backed Agent\n",
    "\n",
    "This example demonstrates how to create an AI agent using LlamaIndex.\n",
    "\n",
    "First install the dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "# pip install \"llama-index-readers-web\" \"llama-index-readers-wikipedia\" \"llama-index-tools-wikipedia\" \"llama-index-embeddings-azure-openai\" \"llama-index-llms-azure-openai\" \"llama-index\" \"azure-identity\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's import the modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import List, Optional\n",
    "\n",
    "from autogen_core import AgentId, MessageContext, RoutedAgent, SingleThreadedAgentRuntime, message_handler\n",
    "from azure.identity import DefaultAzureCredential, get_bearer_token_provider\n",
    "from llama_index.core import Settings\n",
    "from llama_index.core.agent import ReActAgent\n",
    "from llama_index.core.agent.runner.base import AgentRunner\n",
    "from llama_index.core.base.llms.types import (\n",
    "    ChatMessage,\n",
    "    MessageRole,\n",
    ")\n",
    "from llama_index.core.chat_engine.types import AgentChatResponse\n",
    "from llama_index.core.memory import ChatSummaryMemoryBuffer\n",
    "from llama_index.core.memory.types import BaseMemory\n",
    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.azure_openai import AzureOpenAI\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.tools.wikipedia import WikipediaToolSpec\n",
    "from pydantic import BaseModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define our message type that will be used to communicate with the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Resource(BaseModel):\n",
    "    content: str\n",
    "    node_id: str\n",
    "    score: Optional[float] = None\n",
    "\n",
    "\n",
    "class Message(BaseModel):\n",
    "    content: str\n",
    "    sources: Optional[List[Resource]] = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the agent using LLamaIndex's API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LlamaIndexAgent(RoutedAgent):\n",
    "    def __init__(self, description: str, llama_index_agent: AgentRunner, memory: BaseMemory | None = None) -> None:\n",
    "        super().__init__(description)\n",
    "\n",
    "        self._llama_index_agent = llama_index_agent\n",
    "        self._memory = memory\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_user_message(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        # retriever history messages from memory!\n",
    "        history_messages: List[ChatMessage] = []\n",
    "\n",
    "        response: AgentChatResponse  # pyright: ignore\n",
    "        if self._memory is not None:\n",
    "            history_messages = self._memory.get(input=message.content)\n",
    "\n",
    "            response = await self._llama_index_agent.achat(message=message.content, history_messages=history_messages)  # pyright: ignore\n",
    "        else:\n",
    "            response = await self._llama_index_agent.achat(message=message.content)  # pyright: ignore\n",
    "\n",
    "        if isinstance(response, AgentChatResponse):\n",
    "            if self._memory is not None:\n",
    "                self._memory.put(ChatMessage(role=MessageRole.USER, content=message.content))\n",
    "                self._memory.put(ChatMessage(role=MessageRole.ASSISTANT, content=response.response))\n",
    "\n",
    "            assert isinstance(response.response, str)\n",
    "\n",
    "            resources: List[Resource] = [\n",
    "                Resource(content=source_node.get_text(), score=source_node.score, node_id=source_node.id_)\n",
    "                for source_node in response.source_nodes\n",
    "            ]\n",
    "\n",
    "            tools: List[Resource] = [\n",
    "                Resource(content=source.content, node_id=source.tool_name) for source in response.sources\n",
    "            ]\n",
    "\n",
    "            resources.extend(tools)\n",
    "            return Message(content=response.response, sources=resources)\n",
    "        else:\n",
    "            return Message(content=\"I'm sorry, I don't have an answer for you.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting up LlamaIndex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# llm = AzureOpenAI(\n",
    "#     deployment_name=os.getenv(\"AZURE_OPENAI_DEPLOYMENT\"),\n",
    "#     temperature=0.0,\n",
    "#     azure_ad_token_provider = get_bearer_token_provider(DefaultAzureCredential()),\n",
    "#     # api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n",
    "#     azure_endpoint=os.getenv(\"AZURE_OPENAI_ENDPOINT\"),\n",
    "#     api_version=os.getenv(\"AZURE_OPENAI_API_VERSION\"),\n",
    "# )\n",
    "llm = OpenAI(\n",
    "    model=\"gpt-4o\",\n",
    "    temperature=0.0,\n",
    "    api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
    ")\n",
    "\n",
    "# embed_model = AzureOpenAIEmbedding(\n",
    "#     deployment_name=os.getenv(\"AZURE_OPENAI_EMBEDDING_MODEL\"),\n",
    "#     temperature=0.0,\n",
    "#     azure_ad_token_provider = get_bearer_token_provider(DefaultAzureCredential()),\n",
    "#     api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n",
    "#     azure_endpoint=os.getenv(\"AZURE_OPENAI_ENDPOINT\"),\n",
    "#     api_version=os.getenv(\"AZURE_OPENAI_API_VERSION\"),\n",
    "# )\n",
    "embed_model = OpenAIEmbedding(\n",
    "    model=\"text-embedding-ada-002\",\n",
    "    api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
    ")\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "wiki_spec = WikipediaToolSpec()\n",
    "wikipedia_tool = wiki_spec.to_tool_list()[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's test the agent. First we need to create an agent runtime and\n",
    "register the agent, by providing the agent's name and a factory function\n",
    "that will create the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "runtime = SingleThreadedAgentRuntime()\n",
    "await LlamaIndexAgent.register(\n",
    "    runtime,\n",
    "    \"chat_agent\",\n",
    "    lambda: LlamaIndexAgent(\n",
    "        description=\"Llama Index Agent\",\n",
    "        llama_index_agent=ReActAgent.from_tools(\n",
    "            tools=[wikipedia_tool],\n",
    "            llm=llm,\n",
    "            max_iterations=8,\n",
    "            memory=ChatSummaryMemoryBuffer(llm=llm, token_limit=16000),\n",
    "            verbose=True,\n",
    "        ),\n",
    "    ),\n",
    ")\n",
    "agent = AgentId(\"chat_agent\", \"default\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start the agent runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "runtime.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Send a direct message to the agent, and print the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Running step 3cbf60cd-9827-4dfe-a3a9-eaff2bed9b75. Step input: What are the best movies from studio Ghibli?\n",
      "\u001b[1;3;38;5;200mThought: The current language of the user is: English. I need to use a tool to help me answer the question.\n",
      "Action: search_data\n",
      "Action Input: {'query': 'best movies from Studio Ghibli'}\n",
      "\u001b[0m\u001b[1;3;34mObservation: This is a list of works (films, television, shorts etc.) by the Japanese animation studio Studio Ghibli.\n",
      "\n",
      "\n",
      "== Works ==\n",
      "\n",
      "\n",
      "=== Feature films ===\n",
      "\n",
      "\n",
      "=== Television ===\n",
      "\n",
      "\n",
      "=== Short films ===\n",
      "\n",
      "These are short films, including those created for television, theatrical release, and the Ghibli Museum. Original video animation releases and music videos (theatrical and television) are also listed in this section.\n",
      "\n",
      "\n",
      "=== Commercials ===\n",
      "\n",
      "\n",
      "=== Video games ===\n",
      "\n",
      "\n",
      "=== Stage productions ===\n",
      "Princess Mononoke (2013)\n",
      "Nausicaä of the Valley of the Wind (2019)\n",
      "Spirited Away (2022)\n",
      "My Neighbour Totoro (2022)\n",
      "\n",
      "\n",
      "=== Other works ===\n",
      "The works listed here consist of works that do not fall into the above categories. All of these films have been released on DVD or Blu-ray in Japan as part of the Ghibli Gakujutsu Library.\n",
      "\n",
      "\n",
      "=== Exhibitions ===\n",
      "A selection of layout designs for animated productions was exhibited in the Studio Ghibli Layout Designs: Understanding the Secrets of Takahata and Miyazaki Animation exhibition tour, which started in the Museum of Contemporary Art Tokyo (July 28, 2008 to September 28, 2008) and subsequently travelled to different museums throughout Japan and Asia, concluding its tour of Japan in the Fukuoka Asian Art Museum (October 12, 2013 to January 26, 2014) and its tour of Asia in the Hong Kong Heritage Museum (May 14, 2014 to August 31, 2014). Between October 4, 2014 and March 1, 2015 the layout designs were exhibited at Art Ludique in Paris. The exhibition catalogues contain annotated reproductions of the displayed artwork.\n",
      "\n",
      "\n",
      "== Related works ==\n",
      "These works were not created by Studio Ghibli, but were produced by a variety of studios and people who went on to form or join Studio Ghibli. This includes members of Topcraft that went on to create Studio Ghibli in 1985; works produced by Toei Animation, TMS Entertainment, Nippon Animation or other studios and featuring involvement by Hayao Miyazaki, Isao Takahata or other Ghibli staffers. The list also includes works created in cooperation with Studio Ghibli.\n",
      "\n",
      "\n",
      "=== Pre-Ghibli ===\n",
      "\n",
      "\n",
      "=== Cooperative works ===\n",
      "\n",
      "\n",
      "=== Distributive works ===\n",
      "These Western animated films (plus one Japanese film) have been distributed by Studio Ghibli, and now through their label, Ghibli Museum Library.\n",
      "\n",
      "\n",
      "=== Contributive works ===\n",
      "Studio Ghibli has made contributions to the following anime series and movies:\n",
      "\n",
      "\n",
      "== Significant achievements ==\n",
      "The highest-grossing film of 1989 in Japan: Kiki's Delivery Service\n",
      "The highest-grossing film of 1991 in Japan: Only Yesterday\n",
      "The highest-grossing film of 1992 in Japan: Porco Rosso\n",
      "The highest-grossing film of 1994 in Japan: Pom Poko\n",
      "The highest-grossing film of 1995 in Japan; the first Japanese film in Dolby Digital: Whisper of the Heart\n",
      "The highest-grossing film of 2002 in Japan: Spirited Away\n",
      "The highest-grossing film of 2008 in Japan: Ponyo\n",
      "The highest-grossing Japanese film of 2010 in Japan: The Secret World of Arrietty\n",
      "The highest-grossing film of 2013 in Japan: The Wind Rises\n",
      "The first Studio Ghibli film to use computer graphics: Pom Poko\n",
      "The first Miyazaki feature to use computer graphics, and the first Studio Ghibli film to use digital coloring; the first animated feature in Japan's history to gross more than 10 billion yen at the box office and the first animated film ever to win a National Academy Award for Best Picture of the Year: Princess Mononoke\n",
      "The first Studio Ghibli film to be shot using a 100% digital process: My Neighbors the Yamadas\n",
      "The first Miyazaki feature to be shot using a 100% digital process; the first film to gross $200 million worldwide before opening in North America; the film to finally overtake Titanic at the Japanese box office, becoming the top-grossing film in the history of Japanese cinema: Spirited Away\n",
      "The first anime and traditionally animated winner of the Academy Award for Best Animated Feature: Spirited Away at the 75th Academy Awards. They would later win this award for a second time with The Boy and the Heron at the 96th Academy Awards, marking the second time a traditionally animated film won the award.\n",
      "\n",
      "\n",
      "== Notes ==\n",
      "\n",
      "\n",
      "== References ==\n",
      "\u001b[0m> Running step 561e3dd3-d98b-4d37-b612-c99387182ee0. Step input: None\n",
      "\u001b[1;3;38;5;200mThought: I can answer without using any more tools. I'll use the user's language to answer.\n",
      "Answer: Studio Ghibli has produced many acclaimed films over the years. Some of the best and most popular movies from Studio Ghibli include:\n",
      "\n",
      "1. **Spirited Away (2001)** - Directed by Hayao Miyazaki, this film won the Academy Award for Best Animated Feature and is one of the highest-grossing films in Japanese history.\n",
      "2. **My Neighbor Totoro (1988)** - Another classic by Hayao Miyazaki, this film is beloved for its heartwarming story and iconic characters.\n",
      "3. **Princess Mononoke (1997)** - This epic fantasy film, also directed by Miyazaki, is known for its complex themes and stunning animation.\n",
      "4. **Howl's Moving Castle (2004)** - Based on the novel by Diana Wynne Jones, this film features a magical story and beautiful animation.\n",
      "5. **Kiki's Delivery Service (1989)** - A charming coming-of-age story about a young witch starting her own delivery service.\n",
      "6. **Grave of the Fireflies (1988)** - Directed by Isao Takahata, this poignant film is a heartbreaking tale of two siblings struggling to survive during World War II.\n",
      "7. **Ponyo (2008)** - A delightful and visually stunning film about a young fish-girl who wants to become human.\n",
      "8. **The Wind Rises (2013)** - A more mature film by Miyazaki, focusing on the life of an aircraft designer during wartime Japan.\n",
      "9. **The Secret World of Arrietty (2010)** - Based on Mary Norton's novel \"The Borrowers,\" this film tells the story of tiny people living secretly in a human house.\n",
      "10. **Whisper of the Heart (1995)** - A touching story about a young girl discovering her passion for writing.\n",
      "\n",
      "These films are celebrated for their storytelling, animation quality, and emotional depth.\n",
      "\u001b[0m"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Studio Ghibli has produced many acclaimed films over the years. Some of the best and most popular movies from Studio Ghibli include:\n",
      "\n",
      "1. **Spirited Away (2001)** - Directed by Hayao Miyazaki, this film won the Academy Award for Best Animated Feature and is one of the highest-grossing films in Japanese history.\n",
      "2. **My Neighbor Totoro (1988)** - Another classic by Hayao Miyazaki, this film is beloved for its heartwarming story and iconic characters.\n",
      "3. **Princess Mononoke (1997)** - This epic fantasy film, also directed by Miyazaki, is known for its complex themes and stunning animation.\n",
      "4. **Howl's Moving Castle (2004)** - Based on the novel by Diana Wynne Jones, this film features a magical story and beautiful animation.\n",
      "5. **Kiki's Delivery Service (1989)** - A charming coming-of-age story about a young witch starting her own delivery service.\n",
      "6. **Grave of the Fireflies (1988)** - Directed by Isao Takahata, this poignant film is a heartbreaking tale of two siblings struggling to survive during World War II.\n",
      "7. **Ponyo (2008)** - A delightful and visually stunning film about a young fish-girl who wants to become human.\n",
      "8. **The Wind Rises (2013)** - A more mature film by Miyazaki, focusing on the life of an aircraft designer during wartime Japan.\n",
      "9. **The Secret World of Arrietty (2010)** - Based on Mary Norton's novel \"The Borrowers,\" this film tells the story of tiny people living secretly in a human house.\n",
      "10. **Whisper of the Heart (1995)** - A touching story about a young girl discovering her passion for writing.\n",
      "\n",
      "These films are celebrated for their storytelling, animation quality, and emotional depth.\n"
     ]
    }
   ],
   "source": [
    "message = Message(content=\"What are the best movies from studio Ghibli?\")\n",
    "response = await runtime.send_message(message, agent)\n",
    "assert isinstance(response, Message)\n",
    "print(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This is a list of works (films, television, shorts etc.) by the Japanese animation studio Studio Ghibli.\n",
      "\n",
      "\n",
      "== Works ==\n",
      "\n",
      "\n",
      "=== Feature films ===\n",
      "\n",
      "\n",
      "=== Television ===\n",
      "\n",
      "\n",
      "=== Short films ===\n",
      "\n",
      "These are short films, including those created for television, theatrical release, and the Ghibli Museum. Original video animation releases and music videos (theatrical and television) are also listed in this section.\n",
      "\n",
      "\n",
      "=== Commercials ===\n",
      "\n",
      "\n",
      "=== Video games ===\n",
      "\n",
      "\n",
      "=== Stage productions ===\n",
      "Princess Mononoke (2013)\n",
      "Nausicaä of the Valley of the Wind (2019)\n",
      "Spirited Away (2022)\n",
      "My Neighbour Totoro (2022)\n",
      "\n",
      "\n",
      "=== Other works ===\n",
      "The works listed here consist of works that do not fall into the above categories. All of these films have been released on DVD or Blu-ray in Japan as part of the Ghibli Gakujutsu Library.\n",
      "\n",
      "\n",
      "=== Exhibitions ===\n",
      "A selection of layout designs for animated productions was exhibited in the Studio Ghibli Layout Designs: Understanding the Secrets of Takahata and Miyazaki Animation exhibition tour, which started in the Museum of Contemporary Art Tokyo (July 28, 2008 to September 28, 2008) and subsequently travelled to different museums throughout Japan and Asia, concluding its tour of Japan in the Fukuoka Asian Art Museum (October 12, 2013 to January 26, 2014) and its tour of Asia in the Hong Kong Heritage Museum (May 14, 2014 to August 31, 2014). Between October 4, 2014 and March 1, 2015 the layout designs were exhibited at Art Ludique in Paris. The exhibition catalogues contain annotated reproductions of the displayed artwork.\n",
      "\n",
      "\n",
      "== Related works ==\n",
      "These works were not created by Studio Ghibli, but were produced by a variety of studios and people who went on to form or join Studio Ghibli. This includes members of Topcraft that went on to create Studio Ghibli in 1985; works produced by Toei Animation, TMS Entertainment, Nippon Animation or other studios and featuring involvement by Hayao Miyazaki, Isao Takahata or other Ghibli staffers. The list also includes works created in cooperation with Studio Ghibli.\n",
      "\n",
      "\n",
      "=== Pre-Ghibli ===\n",
      "\n",
      "\n",
      "=== Cooperative works ===\n",
      "\n",
      "\n",
      "=== Distributive works ===\n",
      "These Western animated films (plus one Japanese film) have been distributed by Studio Ghibli, and now through their label, Ghibli Museum Library.\n",
      "\n",
      "\n",
      "=== Contributive works ===\n",
      "Studio Ghibli has made contributions to the following anime series and movies:\n",
      "\n",
      "\n",
      "== Significant achievements ==\n",
      "The highest-grossing film of 1989 in Japan: Kiki's Delivery Service\n",
      "The highest-grossing film of 1991 in Japan: Only Yesterday\n",
      "The highest-grossing film of 1992 in Japan: Porco Rosso\n",
      "The highest-grossing film of 1994 in Japan: Pom Poko\n",
      "The highest-grossing film of 1995 in Japan; the first Japanese film in Dolby Digital: Whisper of the Heart\n",
      "The highest-grossing film of 2002 in Japan: Spirited Away\n",
      "The highest-grossing film of 2008 in Japan: Ponyo\n",
      "The highest-grossing Japanese film of 2010 in Japan: The Secret World of Arrietty\n",
      "The highest-grossing film of 2013 in Japan: The Wind Rises\n",
      "The first Studio Ghibli film to use computer graphics: Pom Poko\n",
      "The first Miyazaki feature to use computer graphics, and the first Studio Ghibli film to use digital coloring; the first animated feature in Japan's history to gross more than 10 billion yen at the box office and the first animated film ever to win a National Academy Award for Best Picture of the Year: Princess Mononoke\n",
      "The first Studio Ghibli film to be shot using a 100% digital process: My Neighbors the Yamadas\n",
      "The first Miyazaki feature to be shot using a 100% digital process; the first film to gross $200 million worldwide before opening in North America; the film to finally overtake Titanic at the Japanese box office, becoming the top-grossing film in the history of Japanese cinema: Spirited Away\n",
      "The first anime and traditionally animated winner of the Academy Award for Best Animated Feature: Spirited Away at the 75th Academy Awards. They would later win this award for a second time with The Boy and the Heron at the 96th Academy Awards, marking the second time a traditionally animated film won the award.\n",
      "\n",
      "\n",
      "== Notes ==\n",
      "\n",
      "\n",
      "== References ==\n"
     ]
    }
   ],
   "source": [
    "if response.sources is not None:\n",
    "    for source in response.sources:\n",
    "        print(source.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stop the agent runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "await runtime.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "autogen_core",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
